import sounddevice as sd
import torch
import numpy as np
import random

from dataset import get_files, Audio2ParamsDataset
from simpledaw import SimpleDAW


if __name__ == '__main__':
    
    cwd = 'C:/dev_spa/DMuse/202202b3'
    
    config = {
            'dataset_config': {
                'train_sample_dir': cwd,
                'test_sample_dir': cwd,
                'train_label_dir': cwd,
                'test_label_dir': cwd,
                'compress_rate': 30.,
                'n_spaces': 30,
            },

            'validation_config': {
                'source_sample_rate': 44100.,
                'bpm': 120,
            },
            'batch_size': 1,
            'shuffle': True,
        }
    
    daw = SimpleDAW('C:/VST/64bit/Sylenth1.dll', sample_rate=config['validation_config']['source_sample_rate'])
    
    midi_dir = r"B:\muse_repo\MIDI"
    mid_files = get_files(midi_dir, 'mid')
    daw.load_midi(random.choice(mid_files))
    
    
    dataset_class = Audio2ParamsDataset
    test_dataset = dataset_class(mode='test', **config['dataset_config'])
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=config['batch_size'],
                                                 shuffle=config['shuffle'] if 'shuffle' in config else True)
    network = torch.load(r'B:\codes\ds\DMuse\artear2_1\202202a_model.pth')
    for i, (x_data, y_data) in enumerate(test_loader):
        prediction = network(x_data)
        params = list(torch.argmax(prediction, -1).detach().numpy()[0, ...])
        for i in range(len(params)):
            if params[i] >= config['dataset_config']['n_spaces']:
                params[i] = 1.
            else:
                params[i] /= config['dataset_config']['n_spaces']
        print(params)
        daw.set_params(params)
        audio = daw.render(8.)
        print('正在播放预测参数')
        sd.play(audio, config['validation_config']['source_sample_rate'], blocking=True)
        print('播放结束')
